{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"align-center\">\n",
    "<a href=\"https://oumi.ai/\"><img src=\"https://oumi.ai/docs/en/latest/_static/logo/header_logo.png\" height=\"200\"></a>\n",
    "\n",
    "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n",
    "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n",
    "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n",
    "</div>\n",
    "\n",
    "👋 Welcome to Open Universal Machine Intelligence (Oumi)!\n",
    "\n",
    "🚀 Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](hhttps://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n",
    "\n",
    "🤝 Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n",
    "\n",
    "⭐ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S7bYaH10SgtN"
   },
   "source": [
    "# A Tour of Oumi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QkhNGqE1SgtP"
   },
   "source": [
    "This tutorial will give you a brief overview of Oumi's core functionality. We'll cover:\n",
    "\n",
    "1. Training a model\n",
    "1. Performing model inference\n",
    "1. Evaluating a model against common benchmarks\n",
    "1. Launching jobs\n",
    "1. Customizing datasets and clouds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fHDr11SqSgtP"
   },
   "source": [
    "# 📋 Prerequisites\n",
    "## Oumi Installation\n",
    "\n",
    "First, let's install Oumi. You can find more detailed instructions [here](https://oumi.ai/docs/en/latest/get_started/installation.html). \n",
    "\n",
    "If you have a GPU, you can run the following commands to install Oumi:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n",
      "\u001b[2mUsing Python 3.11.8 environment at: /Users/oussamaelachqar/miniconda3/envs/oumi\u001b[0m\n",
      "\u001b[2mAudited \u001b[1m1 package\u001b[0m \u001b[2min 15ms\u001b[0m\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "%pip install uv -q\n",
    "!uv pip install oumi --no-progress --system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "❗**WARNING:** After the first `pip install`, you may have to restart the notebook for the package updates to take effect (Colab Menu: `Runtime` -> `Restart Session`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JPmWKRVCSgtP"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "tutorial_dir = \"tour_tutorial\"\n",
    "\n",
    "Path(tutorial_dir).mkdir(parents=True, exist_ok=True)\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"  # Disable warnings from HF."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m8PTJuc4SgtQ"
   },
   "source": [
    "# ⚒️ Training a Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b2_HamuySgtQ"
   },
   "source": [
    "Oumi supports training both custom and out-of-the-box models. Want to try out a model on HuggingFace? You can do that. Want to train your own custom Pytorch model? No problem.\n",
    "\n",
    "## A Quick Demo\n",
    "\n",
    "Let's try training a pre-existing model on HuggingFace. We'll use SmolLM2 135M as it's small and trains quickly.\n",
    "\n",
    "Oumi uses [training configuration files](https://oumi.ai/docs/en/latest/api/oumi.core.configs.html#oumi.core.configs.TrainingConfig) to specify training parameters. We've already created a training config for SmolLM2 — let's give it a try!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l2SQ9fZiSgtQ"
   },
   "outputs": [],
   "source": [
    "yaml_content = f\"\"\"\n",
    "model:\n",
    "  model_name: \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  trust_remote_code: True\n",
    "\n",
    "data:\n",
    "  train:\n",
    "    datasets:\n",
    "      - dataset_name: \"yahma/alpaca-cleaned\"\n",
    "    target_col: \"prompt\"\n",
    "\n",
    "training:\n",
    "  trainer_type: \"TRL_SFT\"\n",
    "  per_device_train_batch_size: 2\n",
    "  max_steps: 10 # Quick \"mini\" training, for demo purposes only.\n",
    "  run_name: \"smollm2_135m_sft\"\n",
    "  output_dir: \"{tutorial_dir}/output\"\n",
    "\"\"\"\n",
    "\n",
    "with open(f\"{tutorial_dir}/train.yaml\", \"w\") as f:\n",
    "    f.write(yaml_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2GpQDGG5SgtQ"
   },
   "outputs": [],
   "source": [
    "from oumi.core.configs import TrainingConfig\n",
    "from oumi.train import train\n",
    "\n",
    "config = TrainingConfig.from_yaml(str(Path(tutorial_dir) / \"train.yaml\"))\n",
    "\n",
    "train(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1XfsWKEFSgtR"
   },
   "source": [
    "Congratulations, you've trained your first model using Oumi!\n",
    "\n",
    "You can also train your own custom Pytorch model. We cover that in depth in our [Finetuning Tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Finetuning%20Tutorial.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PdMUxYZcSgtR"
   },
   "source": [
    "# 🧠 Model Inference\n",
    "\n",
    "Now that you've trained a model, let's run inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z4bx6ibXSgtR"
   },
   "outputs": [],
   "source": [
    "yaml_content = f\"\"\"\n",
    "model:\n",
    "  model_name: \"{tutorial_dir}/output\"\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "\n",
    "generation:\n",
    "  max_new_tokens: 128\n",
    "  batch_size: 1\n",
    "\"\"\"\n",
    "\n",
    "with open(f\"{tutorial_dir}/infer.yaml\", \"w\") as f:\n",
    "    f.write(yaml_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XaxxRD1sSgtR"
   },
   "outputs": [],
   "source": [
    "from oumi.core.configs import InferenceConfig\n",
    "from oumi.infer import infer\n",
    "\n",
    "config = InferenceConfig.from_yaml(str(Path(tutorial_dir) / \"infer.yaml\"))\n",
    "\n",
    "input_text = (\n",
    "    \"Remember that we didn't train for long, so the results might not be great.\"\n",
    ")\n",
    "\n",
    "results = infer(config=config, inputs=[input_text])\n",
    "\n",
    "print(results[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xmlA3ynCSgtR"
   },
   "source": [
    "We can also run inference using the pretrained model by slightly tweaking our config:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Wbo2QAxoSgtR"
   },
   "outputs": [],
   "source": [
    "base_model_config = InferenceConfig.from_yaml(str(Path(tutorial_dir) / \"infer.yaml\"))\n",
    "base_model_config.model.model_name = \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n",
    "\n",
    "input_text = \"Input for the pretrained model: What is your name? \"\n",
    "\n",
    "results = infer(config=base_model_config, inputs=[input_text])\n",
    "\n",
    "print(results[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HnxsJBCTSgtR"
   },
   "source": [
    "# 📊 Evaluating a Model against Common Benchmarks\n",
    "\n",
    "You can use Oumi to evaluate pretrained and tuned models against standard benchmarks. For example, let's evaluate our tuned model against `Hellaswag`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tArOVDzVSgtR"
   },
   "outputs": [],
   "source": [
    "yaml_content = f\"\"\"\n",
    "model:\n",
    "  model_name: \"{tutorial_dir}/output\"\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "\n",
    "tasks:\n",
    "  - evaluation_platform: lm_harness\n",
    "    task_name: mmlu_college_computer_science\n",
    "\n",
    "generation:\n",
    "  batch_size: null # This will let LM HARNESS find the maximum possible batch size.\n",
    "output_dir: \"{tutorial_dir}/output/evaluation\"\n",
    "\"\"\"\n",
    "\n",
    "with open(f\"{tutorial_dir}/eval.yaml\", \"w\") as f:\n",
    "    f.write(yaml_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "e-tcwtMhSgtR"
   },
   "outputs": [],
   "source": [
    "from oumi.core.configs import EvaluationConfig\n",
    "from oumi.evaluate import evaluate\n",
    "\n",
    "eval_config = EvaluationConfig.from_yaml(str(Path(tutorial_dir) / \"eval.yaml\"))\n",
    "\n",
    "# Uncomment the following line to run evals against the V1 HuggingFace Leaderboard.\n",
    "# This may take a while.\n",
    "# eval_config.data.datasets[0].dataset_name = \"huggingface_leaderboard_v1\"\n",
    "\n",
    "evaluate(eval_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZUnIyaZfSgtS"
   },
   "source": [
    "# ☁️ Launching Jobs\n",
    "\n",
    "Oftentimes you'll need to run various tasks (training, evaluation, etc.) on remote hardware that's better suited for the task. Oumi can handle this for you by launching jobs on various compute clusters. For more information about running jobs, see our [Running Jobs Remotely tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Running%20Jobs%20Remotely.ipynb). For running jobs on custom clusters, see our [Launching Jobs on Custom Clusters tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Launching%20Jobs%20on%20Custom%20Clusters.ipynb).\n",
    "\n",
    "\n",
    "Today, Oumi supports running jobs on several cloud provider platforms.\n",
    "\n",
    "For the latest list, we can run the `which_clouds` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MJipzQl4SgtS"
   },
   "outputs": [],
   "source": [
    "import oumi.launcher as launcher\n",
    "\n",
    "print(\"Supported Clouds in Oumi:\")\n",
    "for cloud in launcher.which_clouds():\n",
    "    print(cloud)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5VdZqPtpSgtS"
   },
   "source": [
    "Let's run a simple \"Hello World\" job locally to demonstrate how to use the Oumi job launcher. This job will echo `Hello World`, then run the same training job executed above. Running this job on a cloud provider like GCP simply involves changing the `cloud` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "15_fKEiNSgtS"
   },
   "outputs": [],
   "source": [
    "yaml_content = f\"\"\"\n",
    "name: hello-world\n",
    "resources:\n",
    "  cloud: local\n",
    "\n",
    "working_dir: .\n",
    "\n",
    "envs:\n",
    "  TEST_ENV_VARIABLE: '\"Hello, World!\"'\n",
    "  OUMI_LOGGING_DIR: \"{tutorial_dir}/logs\"\n",
    "\n",
    "run: |\n",
    "  echo \"$TEST_ENV_VARIABLE\"\n",
    "  oumi train -c {tutorial_dir}/train.yaml\n",
    "\"\"\"\n",
    "\n",
    "with open(f\"{tutorial_dir}/job.yaml\", \"w\") as f:\n",
    "    f.write(yaml_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VA-n5iUASgtS"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "job_config = launcher.JobConfig.from_yaml(str(Path(tutorial_dir) / \"job.yaml\"))\n",
    "cluster, job_status = launcher.up(job_config, cluster_name=None)\n",
    "\n",
    "while job_status and not job_status.done:\n",
    "    print(\"Job is running...\")\n",
    "    time.sleep(15)\n",
    "    job_status = cluster.get_job(job_status.id)\n",
    "print(\"Job is done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DuVVlpbOSgtS"
   },
   "source": [
    "The job created logs under our tutorial directory. Let's take a look at the directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-BGuSP6XTVKh"
   },
   "outputs": [],
   "source": [
    "logs_dir = f\"{tutorial_dir}/logs\"\n",
    "os.listdir(logs_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N_iz4i4ATXuB"
   },
   "source": [
    "Now let's parse the logfiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RbhUaUchSgtT"
   },
   "outputs": [],
   "source": [
    "for log_file in Path(logs_dir).iterdir():\n",
    "    print(f\"Log file: {log_file}\")\n",
    "    with open(log_file) as f:\n",
    "        print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nnWJ_XgiSgtT"
   },
   "source": [
    "# ⚙️ Customizing Datasets and Clusters \n",
    "\n",
    "Oumi offers rich customization that allows users to build custom solutions on top of our existing building blocks. Several of Oumi's primary resources (Datasets, Clouds, etc.) leverage the Oumi Registry when invoked.\n",
    "\n",
    "This registry allows users to build custom classes that function as drop-in replacements for core functionality.\n",
    "\n",
    "For more details on registering custom datasets, see the [tutorial here](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Datasets%20Tutorial.ipynb).\n",
    "\n",
    "For a tutorial on writing a custom cloud/cluster for running jobs, see the [tutorial here](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Launching%20Jobs%20on%20Custom%20Clusters.ipynb).\n",
    "\n",
    "You can find further information about the required registry decorators [here](https://oumi.ai/docs/en/latest/api/oumi.core.registry.html#oumi.core.registry.register_cloud_builder)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KcM7KZ2OSgtT"
   },
   "source": [
    "# 🧭 What's Next?\n",
    "\n",
    "Now that you've completed the basic tour, you're ready to tackle the other [notebook guides & tutorials](https://oumi.ai/docs/en/latest/get_started/tutorials.html).\n",
    "\n",
    "If you have not already, make sure to take a look at the [Quickstart](https://oumi.ai/docs/en/latest/get_started/quickstart.html) for an overview of our CLI."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "oumi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
